from common_import import *
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.model_selection import train_test_split
from common_import import *
import lightgbm as lgb


def cal_importance_LGBMRegressor(X_train, y_train, X_test):
    # 使用LightGBM训练模型
    model = lgb.LGBMRegressor(
        objective="regression",
        n_estimators=100,
        random_state=40,
        learning_rate=0.05,
        max_depth=15,
    )
    # 拟合模型
    model.fit(X_train, y_train)

    # 使用训练好的模型对测试集进行预测
    y_pred = model.predict(X_test)

    # 返回预测结果
    return y_pred


if __name__ == "__main__":
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    era_activity = pd.read_csv("data/ER_activity_training.csv")
    X = molecular_descriptor.drop(columns=["SMILES"])
    y = era_activity["pIC50"]
    filter_X = X[constants.feature_20]
    X_test = pd.read_csv("data/Molecular_Descriptor_test.csv").drop(columns=["SMILES"])
    filter_Xtest = X_test[constants.feature_20]
    result = cal_importance_LGBMRegressor(filter_X, y, filter_Xtest)
    print(result)
